from scipy import signal
import numpy as np
import matplotlib.pyplot as plt

def CutSignalBiside(wavSignal, fs):
    leftOff = int(0.3 * fs)
    rightOff = int(0.2 * fs)
    wavLen = len(wavSignal)
    wavSignal = wavSignal[leftOff:wavLen-rightOff]
    return wavSignal

def CutSTFTBiside(Zxx):
    # leftOff = 32
    leftOff = 20
    rightOff = 20
    timeLen = Zxx.shape[1]
    Zxx = Zxx[:,leftOff:timeLen-rightOff]
    return Zxx

def CondSignal(var, minTh, winLen):

    varLen =  len(var)
    leftIndex = 0
    rightIndex = varLen-1
    for i in range(varLen - winLen):
        if (var[i+1]<var[i]):
            leftIndex = i+1
            continue
        temp = var[i:i+winLen]
        judge = (temp > minTh)
        if True in judge:
            leftIndex = i
            break

    for i in range(varLen-1, winLen, -1):
        if (var[i-1]<var[i]):
            rightIndex = i - 1
            continue
        temp = var[i-winLen:i]
        judge = (temp > minTh)
        if True in judge:
            rightIndex = i
            break
    return leftIndex, rightIndex

def SplitWordIndex(var, noiseThresh, peakThresh, countThresh, winLen):
    data = []
    count = 0
    lastIndex = 0
    for i in range(len(var)-winLen):
        temp = var[i:i+winLen]
        temp = (temp > noiseThresh)
        # print(temp)
        if True not in temp:
            count += 1
        else:
            count = 0

        if count >= countThresh:
            # print(count)
            active = var[lastIndex:i]
            peak = signal.find_peaks(active)
            if len(peak[0]) > 0:
                judge = (active[peak[0]] > peakThresh)
            else:
                judge = np.array([False])

            if True in judge:
                # active = var[lastIndex:i+winLen]
                # plt.plot(active)
                # plt.show()
                data.append((lastIndex, i+winLen))
                lastIndex = i+winLen
            count = 0

    return data

def SplitWord(Zxx, var, noiseThresh, peakThresh, countThresh, winLen):
    data = []
    indexData = SplitWordIndex(var, noiseThresh, peakThresh, countThresh, winLen)
    print(indexData)
    for index in indexData:
        # temp = var[index[0]:index[1]]
        temp = Zxx[:, index[0]:index[1]]
        data.append(temp)
    return data

